# 提出问题--->文档加载--->文档分割--->embeddings词向量--->存储至向量数据库--->向量数据库根据问题检索数据--->生成含Prompt的问题--->将问题传递给LLM--->LLM给出答案

# 文档分割 --- Langchain提供多种文档分割方式，区别在怎么确定块与块之间的边界、块由哪些字符/token组成、以及如何测量块大小

import os
# 导入文本分割器
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
# 导入Token分割器
from langchain.text_splitter import TokenTextSplitter
# Notion加载器
from langchain.document_loaders import NotionDirectoryLoader
# Markdown分割器
from langchain.text_splitter import MarkdownHeaderTextSplitter

api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"


# 基于字符的分割方式
def split_char():
    '''
    * RecursiveCharacterTextSplitter 递归字符文本分割
    * CharacterTextSplitter 字符分割
    RecursiveCharacterTextSplitter 将按不同的字符递归地分割(按照这个优先级["\n\n", "\n", " ", ""])，这样就能尽量把所有和语义相关的内容尽可能长时间地保留在同一位置
    RecursiveCharacterTextSplitter需要关注的是4个参数：

    * separators - 分隔符字符串数组
    * chunk_size - 每个文档的字符数量限制
    * chunk_overlap - 两份文档重叠区域的长度
    * length_function - 长度计算函数
    CharacterTextSplitter 字符分割
    分隔符的参数是单个的字符串。这就会使得RecursiveCharacterTextSplitter比CharacterTextSplitter对文档切割得更加碎片化
    '''
    # 初始化参数
    chunk_size = 26  # 设置块大小（每块里面的字符数量）
    chunk_overlap = 4  # 设置块重叠大小（与其他块重复的字符数量）
    # 初始化文本分割器
    r_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap
    )
    r_splitter_2 = RecursiveCharacterTextSplitter(
        chunk_size=150,
        chunk_overlap=0,
        separators=["\n\n", "\n","(?<=\. )", " ", ""]
    )
    c_splitter = CharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separator=' '  # 可以看到字符分割器没有分割这个文本，因为字符文本分割器默认以换行符为分隔符，因此需要设置空格为分隔符。
    )
    # 测试文本
    text2 = "abcdefghijklmnopqrstuvwxyzabcdefg"
    # 使用递归字符文本分割器
    result1 = r_splitter.split_text(text2)
    print(result1)

    # 测试文本
    some_text = """在编写文档时，作者将使用文档结构对内容进行分组。 \
        这可以向读者传达哪些想法是相关的。 例如，密切相关的想法\
        是在句子中。 类似的想法在段落中。 段落构成文档。 \n\n\
        段落通常用一个或两个回车符分隔。 \
        回车符是您在该字符串中看到的嵌入的“反斜杠 n”。 \
        句子末尾有一个句号，但也有一个空格。\
        并且单词之间用空格分隔"""
    # 使用递归字符文本分割器
    result2 = r_splitter_2.split_text(some_text)
    print(result2)

    # 测试文本
    text3 = "a b c d e f g h i j k l m n o p q r s t u v w x y z"
    # 使用字符文本分割器
    result3 = c_splitter.split_text(text3)
    print(result3)


# 基于token的分割方式（很多LLM的上下文窗口长度限制是按照Token来计数的。因此，以LLM的视角，按照Token对文本进行分隔，通常可以得到更好的结果。通过一个实例理解基于字符分割和基于Token分割的区别）
def split_token():
    text_splitter = TokenTextSplitter(
        chunk_size=1,  # 将块大小设为1，块重叠大小设为0，相当于将任意字符串分割成了单个Token组成的列
        chunk_overlap=0
    )
    text1 = "foo bar bazzyfoo"
    result1 = text_splitter.split_text(text1)
    print(result1)
    text2 = "test basecese"
    result2 = text_splitter.split_text(text2)
    print(result2)


# 基于Markdown标题的分割（分块的目的是把具有上下文的文本放在一起，我们可以通过使用指定分隔符来进行分隔，但有些类型的文档（例如 Markdown）本身就具有可用于分割的结构（如标题）。Markdown标题文本分割器会根据标题或子标题来分割一个Markdown文档，并将标题作为元数据添加到每个块中）
def split_markdown():
    # 定义一个Markdonw文本
    markdown_document = """# Title\n\n \
    ## 第一章\n\n \
    李白乘舟将欲行\n\n 忽然岸上踏歌声\n\n \
    ### Section \n\n \
    桃花潭水深千尺 \n\n 
    ## 第二章\n\n \
    不及汪伦送我情"""
    # 定义想要分割的标题列表和名称
    headers_to_split_on = [
        ("#", "Header 1"),
        ("##", "Header 2"),
        ("###", "Header 3"),
    ]

    # 初始化Markdown标题文本分割器，分割Markdown文档
    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=headers_to_split_on
    )
    md_header_splits = markdown_splitter.split_text(markdown_document)
    print(md_header_splits[0])
    print(md_header_splits[1])


# 读取文档并切割（示例）
def split_markdown_example():
    loader = NotionDirectoryLoader("docs/Notion_DB")
    docs = loader.load()
    txt = ' '.join([d.page_content for d in docs])  # 拼接文档
    headers_to_split_on = [
        ("#", "Header 1"),
        ("##", "Header 2"),
    ]
    # 加载文档分割器
    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=headers_to_split_on
    )
    md_header_splits = markdown_splitter.split_text(txt)  # 分割文本内容
    print(md_header_splits[0])  # 分割结果


if __name__ == '__main__':
    # split_char()
    # split_token()
    # split_markdown()
    split_markdown_example()
